package main;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintStream;
import java.io.StringReader;


import edu.washington.cs.knowitall.extractor.ChunkedArgumentExtractor;
import edu.washington.cs.knowitall.extractor.ExtractorUnion;
import edu.washington.cs.knowitall.extractor.RegexRelationExtractor;
import edu.washington.cs.knowitall.extractor.RelationFirstNpChunkExtractor;
import edu.washington.cs.knowitall.extractor.ChunkedArgumentExtractor.Mode;
import edu.washington.cs.knowitall.extractor.mapper.ReVerbArgument1Mappers;
import edu.washington.cs.knowitall.extractor.mapper.ReVerbArgument2Mappers;
import edu.washington.cs.knowitall.nlp.ChunkedSentence;
import edu.washington.cs.knowitall.nlp.ChunkedSentenceReader;
import edu.washington.cs.knowitall.nlp.extraction.ChunkedBinaryExtraction;
import edu.washington.cs.knowitall.nlp.extraction.ChunkedExtraction;
import edu.washington.cs.knowitall.util.DefaultObjects;


public class ReVerbExtractor extends RelationFirstNpChunkExtractor {

    /**
     * Definition of the "verb" of the relation pattern.
     */
    public static final String VERB = 
        "POS_RB? [POS_MD POS_VB POS_VBD POS_VBP POS_VBZ POS_VBG POS_VBN] POS_RP? POS_RB?";

    /**
     * Definition of the "any word" part of the relation pattern.
     */
    public static final String WORD = 
        "[POS_$ POS_PRP$ POS_CD POS_DT POS_JJ POS_JJS POS_JJR POS_NN POS_NNS POS_NNP POS_NNPS POS_POS POS_PRP POS_RB POS_RBR POS_RBS POS_VBN POS_VBG]";

    /**
     * Definition of the "preposition" part of the relation pattern.
     */
    public static final String PREP = "POS_RB? [POS_IN POS_TO POS_RP] POS_RB?";


    // The long pattern is (V(W*P)?)+
    public static final String LONG_RELATION_PATTERN = String.format("(%s (%s* (%s)+)?)+", VERB, WORD, PREP);
    
    // The short pattern is (VP?)+
    public static final String SHORT_RELATION_PATTERN = String.format("(%s (%s)?)+", VERB, PREP);

	private static final String NATURAL_DISASTER = "NaturalDisaster";

    /**
     * Constructs a new <code>ReVerbExtractor</code> using the default relation pattern,
     * relation mappers, and argument mappers.
     */
    public ReVerbExtractor() throws IOException {
        initializeRelationExtractor();
        initializeArgumentExtractors();
    }

    public ReVerbExtractor(String freqFile) throws IOException {
		
    	initializeRelationExtractor(freqFile);
        initializeArgumentExtractors();
    	
	}

	private void initializeRelationExtractor(String freqFile) throws IOException {
    	ExtractorUnion<ChunkedSentence, ChunkedExtraction> relExtractor = new ExtractorUnion<ChunkedSentence, ChunkedExtraction>();
    	relExtractor.addExtractor(new RegexRelationExtractor(SHORT_RELATION_PATTERN));
    	relExtractor.addExtractor(new RegexRelationExtractor(LONG_RELATION_PATTERN));
        relExtractor.addMapper(new ReVerbRelationMappers(freqFile));
        setRelationExtractor(relExtractor);

	}

	private void initializeArgumentExtractors() {
    	ChunkedArgumentExtractor arg1Extr = new ChunkedArgumentExtractor(ChunkedArgumentExtractor.Mode.LEFT);
        arg1Extr.addMapper(new ReVerbArgument1Mappers());
        setArgument1Extractor(arg1Extr);

        ChunkedArgumentExtractor arg2Extr = new ChunkedArgumentExtractor(ChunkedArgumentExtractor.Mode.RIGHT);
        arg2Extr.addMapper(new ReVerbArgument2Mappers());
        setArgument2Extractor(arg2Extr);
    }

    private void initializeRelationExtractor() throws IOException {
    	ExtractorUnion<ChunkedSentence, ChunkedExtraction> relExtractor = new ExtractorUnion<ChunkedSentence, ChunkedExtraction>();
    	relExtractor.addExtractor(new RegexRelationExtractor(SHORT_RELATION_PATTERN));
    	relExtractor.addExtractor(new RegexRelationExtractor(LONG_RELATION_PATTERN));
        relExtractor.addMapper(new ReVerbRelationMappers());
        setRelationExtractor(relExtractor);
    }
    
    /**
     * Extracts from the given text using the default sentence reader returned by <code>DefaultObjects.getDefaultSentenceReader()</code>.
     * @param text
     * @return an <code>Iterable</code> over the extractions
     * @throws IOException
     */
    public Iterable<ChunkedBinaryExtraction> extractFromString(String text) throws IOException {
    	StringReader in = new StringReader(text);
    	return extract(DefaultObjects.getDefaultSentenceReader(in).getSentences());
    }
    
    /**
     * Extracts from the given html using the default sentence reader returned by <code>DefaultObjects.getDefaultSentenceReaderHtml()</code>.
     * @param html
     * @return an <code>Iterable</code> over the extractions
     * @throws IOException
     */
    public Iterable<ChunkedBinaryExtraction> extractFromHtml(String html) throws IOException {
    	StringReader in = new StringReader(html);
    	return extract(DefaultObjects.getDefaultSentenceReaderHtml(in).getSentences());
    }

    /**
     * Runs the extractor on either standard input, or the given file. Uses the object returned by
     * the <code>DefaultObjects.getDefaultSentenceReaderHtml</code> method to read <code>NpChunkedSentence</code>
     * objects. Prints each sentence (prefixed by "sentence" and then a tab), followed by the extractions in the
     * form "extraction", arg1, relation, and arg2, separated by tabs. 
     * @param args
     * @throws Exception
     */
    public static void main(String[] args) throws Exception {

		String[] options = {"_acc_freq_rel.train.txt.gz", "_freq_rel.train.txt.gz"};
		
		String[] Datasets = {"AFFILIATE-PARTNER","ASSOCIATE","AT","BASED-IN","CITIZEN-OF","CLIENT","FOUNDER","GENERAL-STAFF","GRANDPARENT",
				"LOCATED", "MANAGEMENT","MEMBER","NEAR","OTHER","OTHER-PERSONAL","OTHER-PROFESSIONAL","OTHER-RELATIVE","OWNER","PARENT",
				"PART","PART-OF","RELATIVE-LOCATION","RESIDENCE","ROLE","SIBLING","SOC","SPOUSE","SUBSIDIARY","NaturalDisaster"};

		for (int d = 0; d < Datasets.length; d++) {

			int[] split = Chunking.getSplits(Datasets[d]);
			
			for (int i = 0; i < split.length; i++) {
				
				for (int opt = 0; opt < options.length; opt++) {
					
					args = new String[1];
					
					args[0] = Chunking.getDirectory(Datasets[d]) + i + "/testing_list";
					
					System.setOut(new PrintStream(Chunking.getDirectory(Datasets[d]) + i + "/reverb."+opt+".extract"));
					
					String folder = Chunking.getDirectory(Datasets[d]) + i + "/testingDirectory/";
					
			        BufferedReader reader;
			        if (args.length == 0) {
			            reader = new BufferedReader(new InputStreamReader(System.in));
			        } else {
			            reader = new BufferedReader(new FileReader(args[0]));
			        }
			        
			        int sentenceCount = 0;
			        int extractionCount = 0;
			        
			        System.err.print("Initializing extractor...");
			        
			        String freq = Chunking.getDirectory(Datasets[d])+i+"/"+Datasets[d]+options[opt];
			        
			        ReVerbExtractor extractor = new ReVerbExtractor(freq); //need my files.
			        
			        System.err.println("Done.");
			        
			        System.err.print("Initializing confidence function...");
			        
			        String confFunctionModel = Chunking.getDirectory(Datasets[d]) + i + "/" + Datasets[d] + ".reverb.model";
			        
			        ReVerbConfFunction scoreFunc = new ReVerbConfFunction(confFunctionModel);
			        
			        System.err.println("Done.");
			        
			        String file = reader.readLine();
			       
			        while (file!=null){
			        
			            System.err.print("Initializing NLP tools...");
			            ChunkedSentenceReader sentReader = DefaultObjects.getDefaultSentenceReader(new BufferedReader(new FileReader(folder + file)));
			            System.err.println("Done.");

			        	
				        for (ChunkedSentence sent : sentReader.getSentences()) {
				        	
				            sentenceCount++;
				            
				            String sentString = sent.getTokensAsString();
				            System.out.println(String.format("sentence\t%s\t%s", sentenceCount, sentString));
				            
				            
				            for (ChunkedBinaryExtraction extr : extractor.extract(sent)) {
				            	
				                double score = scoreFunc.getConf(extr);
				                
				                String arg1 = extr.getArgument1().toString();
				                String rel = extr.getRelation().toString();
				                String arg2 = extr.getArgument2().toString();
				                
				                String extrString = String.format("%s\t%s\t%s\t%s\t%s", sentenceCount, arg1, rel, arg2, score);
				                
				                System.out.println("extraction\t" + extrString);
				                
				                extractionCount++;
				            }
				        }
				        
				        System.err.println(String.format("Got %s extractions from %s sentences.", extractionCount, sentenceCount));
				        
				        file = reader.readLine();
			        
			        }

					
				}
				
			}
			

			
		}
    }

}
    
